"""TCP隧道客户端"""
import asyncio
import json
import logging
from base64 import b64decode, b64encode
import aiohttp
from homeassistant.core import HomeAssistant
from homeassistant.config_entries import ConfigEntry
from homeassistant.helpers.aiohttp_client import async_get_clientsession

from .const import (
    WEBSOCKET_HEARTBEAT_INTERVAL,
    WEBSOCKET_RECONNECT_INTERVAL,
)

_LOGGER = logging.getLogger(__name__)

class TCPTunnelClient:
    """TCP隧道客户端，完全透传TCP数据"""
    
    def __init__(self, hass: HomeAssistant, entry: ConfigEntry, server_url: str, connection_state_callback=None):
        self.hass = hass
        self.entry = entry
        self.server_url = server_url
        self.session = async_get_clientsession(hass)
        self.ws_client = None
        self.heartbeat_task = None
        self.reconnect_task = None
        self.message_handler_task = None
        self.closing = False
        self._connection_lock = asyncio.Lock()
        self.connections = {}  # connId -> (reader, writer)
        self._connection_state_callback = connection_state_callback
        self._connected = False
        self._connecting = False
        
        # 获取配置
        self.username = entry.data["username"]
        self.password = entry.data["password"]
    
    async def _start_reconnect(self):
        """如果未运行则启动重连任务"""
        if not self.closing and not self.reconnect_task:
            self.reconnect_task = asyncio.create_task(self._reconnect())
    
    async def connect(self):
        """连接到隧道服务器"""
        async with self._connection_lock:
            if self._connected or self._connecting:
                return
                
            self._connecting = True
            self._update_connection_state()
                
            # 确保之前的连接已经完全关闭
            await self._cleanup_connection(start_reconnect=False)
            
            try:
                url = f"{self.server_url}?username={self.username}&password={self.password}"
                self.ws_client = await self.session.ws_connect(url)
                self._connected = True
                self._connecting = False
                self._update_connection_state()
                _LOGGER.info("TCP隧道客户端连接成功")
                
                # 启动心跳任务
                if not self.heartbeat_task:
                    self.heartbeat_task = asyncio.create_task(self._heartbeat())
                
                # 启动消息处理任务
                if not self.message_handler_task:
                    self.message_handler_task = asyncio.create_task(self._handle_messages())
                    
            except Exception as err:
                _LOGGER.error("连接到TCP隧道服务器失败: %s", err)
                await self._cleanup_connection(start_reconnect=True)

    async def _cleanup_connection(self, start_reconnect: bool = False):
        """清理当前连接和所有相关任务"""
        self._connected = False
        self._connecting = False
        self._update_connection_state()
        
        # 取消心跳任务
        if self.heartbeat_task:
            self.heartbeat_task.cancel()
            try:
                await self.heartbeat_task
            except asyncio.CancelledError:
                pass
            self.heartbeat_task = None
        
        # 取消消息处理任务
        if self.message_handler_task:
            self.message_handler_task.cancel()
            try:
                await self.message_handler_task
            except asyncio.CancelledError:
                pass
            self.message_handler_task = None
            
        # 关闭WebSocket连接
        if self.ws_client:
            if not self.ws_client.closed:
                await self.ws_client.close()
            self.ws_client = None
            
        # 清理所有TCP连接
        for conn_id in list(self.connections.keys()):
            await self._handle_close(conn_id)

        # 根据参数决定是否启动重连
        if start_reconnect:
            await self._start_reconnect()
    
    async def _reconnect(self):
        """重新连接到TCP隧道服务器"""
        try:
            while not self.closing:
                _LOGGER.info("将在 %d 秒后尝试重新连接", WEBSOCKET_RECONNECT_INTERVAL)
                await asyncio.sleep(WEBSOCKET_RECONNECT_INTERVAL)
                if not self.closing:
                    await self.connect()
                    if self._connected:
                        break
        finally:
            self.reconnect_task = None
    
    async def _heartbeat(self):
        """发送定期心跳消息"""
        try:
            while True:
                if not self._connected or not self.ws_client or self.ws_client.closed:
                    break
                    
                try:
                    await self.send_message({
                        "type": "ping",
                        "timestamp": int(self.hass.loop.time())
                    })
                    await asyncio.sleep(WEBSOCKET_HEARTBEAT_INTERVAL)
                except Exception as err:
                    _LOGGER.error("心跳错误: %s", err)
                    break
        finally:
            # 如果心跳失败，触发清理和重连
            if not self.closing and self._connected:
                await self._cleanup_connection(start_reconnect=True)
    
    async def send_message(self, message):
        """发送消息到服务器"""
        if not self.ws_client or self.ws_client.closed:
            return
            
        try:
            await self.ws_client.send_json(message)
        except Exception as err:
            _LOGGER.error("发送消息失败: %s", err)
    
    async def _handle_messages(self):
        """处理传入的WebSocket消息"""
        if not self.ws_client:
            return
            
        try:
            async for msg in self.ws_client:
                if msg.type == aiohttp.WSMsgType.TEXT:
                    try:
                        data = json.loads(msg.data)
                        await self._process_message(data)
                    except json.JSONDecodeError as err:
                        _LOGGER.error("解码消息失败: %s", err)
                elif msg.type in (aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.ERROR):
                    break
        except Exception as err:
            _LOGGER.error("TCP隧道消息处理错误: %s", err)
        finally:
            if not self.closing:
                await self._cleanup_connection(start_reconnect=True)

    async def _process_message(self, data):
        """处理消息"""
        msg_type = data.get("type")
        
        if msg_type == "connected":
            _LOGGER.info(f"TCP隧道已建立，子域名: {data.get('subdomain')}")
        elif msg_type == "new_connection":
            await self._handle_new_connection(data["connId"])
        elif msg_type == "data":
            await self._handle_data(data["connId"], data["data"])
        elif msg_type == "close":
            await self._handle_close(data["connId"])

    async def _handle_new_connection(self, conn_id):
        """处理新连接"""
        try:
            # 连接到本地HA，设置超时
            reader, writer = await asyncio.wait_for(
                asyncio.open_connection('127.0.0.1', 8123),
                timeout=10.0
            )
            self.connections[conn_id] = (reader, writer)
            _LOGGER.info(f"新连接建立: {conn_id}")
            
            # 启动数据转发任务
            asyncio.create_task(
                self._forward_local_data(conn_id, reader)
            )
            
        except asyncio.TimeoutError:
            _LOGGER.error(f"连接HA超时: {conn_id}")
        except Exception as e:
            _LOGGER.error(f"连接本地HA失败: {e}")

    async def _forward_local_data(self, conn_id, reader):
        """转发本地数据到服务器"""
        try:
            while conn_id in self.connections:
                try:
                    data = await asyncio.wait_for(reader.read(8192), timeout=30.0)
                    if not data:
                        break
                        
                    if self.ws_client and not self.ws_client.closed:
                        await self.ws_client.send_json({
                            "type": "data",
                            "connId": conn_id,
                            "data": b64encode(data).decode()
                        })
                except asyncio.TimeoutError:
                    # 超时继续循环
                    continue
                except Exception as e:
                    _LOGGER.error(f"读取数据错误: {e}")
                    break
        except Exception as e:
            _LOGGER.error(f"转发本地数据错误: {e}")
        finally:
            if conn_id in self.connections:
                _, writer = self.connections[conn_id]
                writer.close()
                try:
                    await writer.wait_closed()
                except:
                    pass
                self.connections.pop(conn_id, None)
                # 通知服务器关闭连接
                if self.ws_client and not self.ws_client.closed:
                    try:
                        await self.ws_client.send_json({
                            "type": "close",
                            "connId": conn_id
                        })
                    except:
                        pass

    async def _handle_data(self, conn_id, data):
        """处理接收到的数据"""
        if conn_id in self.connections:
            _, writer = self.connections[conn_id]
            try:
                if writer.is_closing():
                    return
                decoded_data = b64decode(data)
                writer.write(decoded_data)
                await asyncio.wait_for(writer.drain(), timeout=5.0)
            except asyncio.TimeoutError:
                _LOGGER.error(f"写入数据超时: {conn_id}")
                await self._handle_close(conn_id)
            except Exception as e:
                _LOGGER.error(f"写入数据错误: {e}")
                await self._handle_close(conn_id)

    async def _handle_close(self, conn_id):
        """处理连接关闭"""
        if conn_id in self.connections:
            _, writer = self.connections[conn_id]
            writer.close()
            try:
                await writer.wait_closed()
            except:
                pass
            self.connections.pop(conn_id, None)
            _LOGGER.info(f"关闭本地连接: {conn_id}")

    def _update_connection_state(self):
        """更新连接状态"""
        if self._connection_state_callback:
            self._connection_state_callback(self._connected, self._connecting)
    
    async def disconnect(self):
        """断开与TCP隧道服务器的连接"""
        self.closing = True
        
        # 取消重连任务
        if self.reconnect_task:
            self.reconnect_task.cancel()
            try:
                await self.reconnect_task
            except asyncio.CancelledError:
                pass
            self.reconnect_task = None
            
        # 清理连接
        await self._cleanup_connection(start_reconnect=False)